import torch.nn as nn  # 导入 PyTorch 的神经网络模块

class WakeWordCRNN(nn.Module):  # 定义唤醒词检测模型类，继承自 nn.Module
    def __init__(self, input_dim=39):  # 初始化方法，接收输入特征维度，默认为 39（MFCC 特征）
        super().__init__()  # 调用父类的初始化方法

        # 改进的 CNN 部分
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=(3, 5), padding=(1, 2)),  # 第一层卷积：输入通道 1，输出通道 32，卷积核 (3, 5)，填充 (1, 2)
            nn.BatchNorm2d(32),  # 批归一化层，加速训练并稳定模型
            nn.GELU(),  # 使用 GELU 激活函数，增加非线性表达能力
            nn.MaxPool2d((2, 2)),  # 最大池化层，下采样 (2, 2)

            nn.Conv2d(32, 64, kernel_size=(3, 5), padding=(1, 2)),  # 第二层卷积：输入通道 32，输出通道 64
            nn.BatchNorm2d(64),  # 批归一化层
            nn.GELU(),  # GELU 激活函数
            nn.MaxPool2d((2, 2)),  # 最大池化层

            nn.AdaptiveAvgPool2d((None, 25))  # 自适应平均池化层，动态调整时间维度为固定长度 25
        )

        # 改进的 RNN 部分
        self.gru = nn.GRU(
            input_size=64 * (input_dim // 4),  # 输入大小为 64 * (input_dim // 4)，即特征图展平后的维度
            hidden_size=128,  # 隐藏层大小为 128
            bidirectional=True,  # 使用双向 GRU
            num_layers=2,  # 增加层数至 2 层
            dropout=0.3  # 在 GRU 层之间添加 Dropout，防止过拟合
        )

        # 改进分类头
        self.classifier = nn.Sequential(
            nn.Linear(256, 128),  # 全连接层：输入维度 256（双向 GRU 输出），输出维度 128
            nn.GELU(),  # GELU 激活函数
            nn.Dropout(0.4),  # Dropout 层，随机丢弃 40% 神经元
            nn.Linear(128, 2)  # 输出层：输入维度 128，输出维度 2（唤醒/非唤醒）
        )

    def forward(self, x):  # 前向传播方法
        x = x.unsqueeze(1)  # 增加通道维度，形状变为 (B, 1, F, T)
        x = self.cnn(x)  # 通过 CNN 提取特征，输出形状为 (B, C, F, T)

        # 调整维度
        B, C, F, T = x.size()  # 获取批次大小、通道数、频率维度、时间步长
        x = x.permute(0, 3, 1, 2)  # 调整维度顺序为 (B, T, C, F)
        x = x.reshape(B, T, -1)  # 将特征图展平为 (B, T, C*F)

        # 使用全部时间步输出
        x, _ = self.gru(x)  # 通过双向 GRU，输出形状为 (B, T, 256)
        x = x.mean(dim=1)  # 对时间步求平均，得到形状 (B, 256)

        return self.classifier(x)  # 通过分类头，输出最终预测结果
    #这段代码实现了一个基于 CRNN 的唤醒词检测模型，结合了卷积神经网络（CNN）和循环神经网络（RNN）的优势，能够高效地处理音频特征并进行分类。